package edu.unika.aifb.graphindex.searcher.keyword.model;

import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.SortedMap;
import java.util.TreeMap;

import edu.unika.aifb.graphindex.model.IEntity;
import edu.unika.aifb.graphindex.searcher.keyword.evaluation.TopKQueryEvaluator;
import edu.unika.aifb.graphindex.util.Util;

public class RelevanceModel {
	private Set<IEntity> entities = new HashSet<IEntity>();
	private SortedMap<String, Integer> termFrequencyMap;
	private Map<String, Double> normalizedTermFrequencyMap;
	private Set<String> keywords;
	private int maxTermsCount = 20;
	private Set<String> stopWords;
	private TopKQueryEvaluator topKQueryEvaluator;
	
	public RelevanceModel(Map<KeywordSegment, Collection<KeywordElement>> segment2elements, int maxTermsCount, TopKQueryEvaluator topKQueryEvaluator){
		this.topKQueryEvaluator = topKQueryEvaluator;
		keywords = new HashSet<String>();
		stopWords = Util.readStopWords();
		for(Collection<KeywordElement> keywordElements : segment2elements.values()){
			for(KeywordElement keywordElement : keywordElements){
				if(keywordElement.getType() == KeywordElement.ENTITY){
					for(IEntity entity : keywordElement.entities){
						entities.add(entity);
					}
					
					keywords.addAll(keywordElement.getKeywordSegment().getKeywords());
				}
			}
		}
		this.maxTermsCount = maxTermsCount;
		calculateRelevanceModel();
	}
	
	public RelevanceModel(IEntity entity, Set<String> keywords, int maxTermsCount){
		stopWords = Util.readStopWords();
		this.keywords = keywords;
		entities.add(entity);
		calculateRelevanceModel();
		this.maxTermsCount = maxTermsCount;
	}

	private void calculateRelevanceModel(){
		Map<String, Integer> unsortedTermFrequenciesMap = new HashMap<String, Integer>();

		for(IEntity entity : entities){
			String[] terms = entity.getTerms();
			int[] freqs = entity.getTermFrequencies();
			for(int i = 0 ; i < terms.length; i++){
				if(terms[i].equalsIgnoreCase("null") || (!keywords.contains(terms[i]) && stopWords.contains(terms[i])))
					continue;
				Integer freq = unsortedTermFrequenciesMap.get(terms[i]);
				if(freq == null)
					unsortedTermFrequenciesMap.put(terms[i], freqs[i]);
				else
					unsortedTermFrequenciesMap.put(terms[i], freq + freqs[i]);
			}
			
		}
		ValueComparator<Integer> integerValueComparator = new ValueComparator<Integer>(unsortedTermFrequenciesMap);
		termFrequencyMap = new TreeMap<String, Integer>(integerValueComparator);
		termFrequencyMap.putAll(unsortedTermFrequenciesMap);
		Map<String, Double> unsortedNormalizedFrequencyMap = new HashMap<String, Double>();
		int termsCount = 0;
		for(Integer  termCount : termFrequencyMap.values()){
			termsCount += termCount;
		}
		for(String term : termFrequencyMap.keySet()){
			unsortedNormalizedFrequencyMap.put(term, ((double)termFrequencyMap.get(term)) / termsCount);
		}
		
		integerValueComparator.clear();
		termFrequencyMap.clear();

		ValueComparator<Double> valueComparator = new ValueComparator<Double>(unsortedNormalizedFrequencyMap);
		normalizedTermFrequencyMap = new TreeMap<String, Double>(valueComparator);
		normalizedTermFrequencyMap.putAll(unsortedNormalizedFrequencyMap);
	
		if(entities.size() > 1){
			unsortedNormalizedFrequencyMap = new HashMap<String, Double>();
			int i = 0;
			double termProbSum = 0;
			for(String term : normalizedTermFrequencyMap.keySet()){
//				double collectionTermProb = topKQueryEvaluator.getBackgroundTermProbability(term);
//				if(!keywords.contains(term) && collectionTermProb > 0.0001){
//					continue;
//				}
				double termProb = normalizedTermFrequencyMap.get(term);
				termProbSum += termProb;
				if(keywords.contains(term)){
					termProbSum += 1;
					unsortedNormalizedFrequencyMap.put(term, termProb + 1);
				}else{
					unsortedNormalizedFrequencyMap.put(term, termProb);
				}
				if(i++ == maxTermsCount)
					break;
			}
			for(String keyword : keywords){
				if(!unsortedNormalizedFrequencyMap.containsKey(keyword)){
					double termProb = normalizedTermFrequencyMap.get(keyword);
					termProbSum += (termProb + 1);
					unsortedNormalizedFrequencyMap.put(keyword, termProb + 1);
				}
			}
			
			for(String term : unsortedNormalizedFrequencyMap.keySet()){
				unsortedNormalizedFrequencyMap.put(term, unsortedNormalizedFrequencyMap.get(term) / termProbSum);
			}
			
			valueComparator.clear();
			normalizedTermFrequencyMap.clear();
			
			valueComparator = new ValueComparator<Double>(unsortedNormalizedFrequencyMap);
			normalizedTermFrequencyMap = new TreeMap<String, Double>(new ValueComparator<Double>(unsortedNormalizedFrequencyMap));
			normalizedTermFrequencyMap.putAll(unsortedNormalizedFrequencyMap);
		}
		
	}
	
	public Map<String, Integer> getTermFrequencies(){
		return termFrequencyMap;
	}
	
	public Map<String, Double> getNormalizedTermFrequencies(){
		return normalizedTermFrequencyMap;
	}
	
	public Set<String> getKeywords() {
		return keywords;
	}



	class ValueComparator<T> implements Comparator<String>{
		Map<String, T> base;
		public ValueComparator(Map<String, T> base){
			this.base = base;
		}
		
		@SuppressWarnings("unchecked")
		public int compare(String a, String b) {
			if(a.equalsIgnoreCase(b))
				return 0;
			T aValue = base.get(a);
			T bValue = base.get(b);
			if(aValue == null || bValue == null || ((Comparable<T>)aValue).compareTo(bValue) == 0){
				return a.compareTo(b);
			}
			
			return ((Comparable<T>)bValue).compareTo(aValue);
		}
		
		public void clear(){
			base.clear();
		}
	}

}
